# scepter_motion.py - MoveIt Task Constructor 机械臂运动控制
from math import pi
import time

from std_msgs.msg import Header
from geometry_msgs.msg import Vector3Stamped, Vector3
from moveit.task_constructor import core, stages

import rclcpp


def main():
    rclcpp.init()
    node = rclcpp.Node("scepter_mtc_example")
    
    # 设置规划组
    arm_group = "arm_group"
    gripper_group = "gripper_group"

    # 创建任务并加载机器人模型
    task = core.Task()
    task.name = "scepter_pick_and_place"
    
    try:
        task.loadRobotModel(node)
        print("Robot model loaded successfully")
    except Exception as e:
        print(f"Failed to load robot model: {e}")
        return

    # 创建规划器 - 使用关节插值和笛卡尔路径规划器
    cartesian_planner = core.CartesianPath()
    joint_planner = core.JointInterpolationPlanner()
    
    print("Planners created successfully")
    
    # 从当前状态开始
    current_state = stages.CurrentState("current state")
    task.add(current_state)

    # 1. 移动到准备位置
    move_to_ready = stages.MoveTo("move to ready", joint_planner)
    move_to_ready.group = arm_group
    move_to_ready.setGoal("ready_pose")
    task.add(move_to_ready)

    # 2. 打开夹爪
    open_gripper = stages.MoveTo("open gripper", joint_planner)
    open_gripper.group = gripper_group
    open_gripper.setGoal("open_pose")
    task.add(open_gripper)

    # 3. 小幅笛卡尔移动 (模拟接近物体)
    move_approach = stages.MoveRelative("approach simulation", cartesian_planner)
    move_approach.group = arm_group
    header = Header(frame_id="world")
    move_approach.setDirection(
        Vector3Stamped(header=header, vector=Vector3(x=0.0, y=0.0, z=-0.05)))  # 减小移动距离
    task.add(move_approach)

    # 4. 关闭夹爪
    close_gripper = stages.MoveTo("close gripper", joint_planner)
    close_gripper.group = gripper_group
    close_gripper.setGoal("close_pose")
    task.add(close_gripper)

    # 5. 小幅向上移动 (模拟提起物体)
    lift_object = stages.MoveRelative("lift simulation", cartesian_planner)
    lift_object.group = arm_group
    lift_object.setDirection(
        Vector3Stamped(header=header, vector=Vector3(x=0.0, y=0.0, z=0.05)))   # 减小移动距离
    task.add(lift_object)

    # 6. 移动到完成位置
    move_to_finish = stages.MoveTo("move to finish", joint_planner)
    move_to_finish.group = arm_group
    move_to_finish.setGoal("finish_pose")  # 使用 SRDF 中定义的姿态
    task.add(move_to_finish)

    # 7. 打开夹爪释放物体
    release_object = stages.MoveTo("release object", joint_planner)
    release_object.group = gripper_group
    release_object.setGoal("open_pose")
    task.add(release_object)

    # 8. 回到初始位置
    move_home = stages.MoveTo("move home", joint_planner)
    move_home.group = arm_group
    move_home.setGoal("stand_pose")
    task.add(move_home)

    # 执行规划
    if task.plan():
        print("Planning succeeded")
        task.publish(task.solutions[0])
    else:
        print("Planning failed")

    # 等待一段时间以便观察结果
    time.sleep(10)


if __name__ == "__main__":
    main()
